import torch
import numpy as np
import pandas as pd
from config import MODEL_CONFIG, TRAIN_CONFIG, DATA_CONFIG, SAVE_CONFIG
from models.lstm_model import create_model
from utils import load_model, create_sequences
from sklearn.preprocessing import StandardScaler

def predict_new_data(model_path, new_data, sequence_length, device):
    """使用训练好的模型进行预测"""
    
    # 加载模型
    model = load_model(create_model, model_path, device)
    model.eval()
    
    # 确保数据是numpy数组
    if isinstance(new_data, pd.DataFrame):
        new_data = new_data.values
    elif isinstance(new_data, list):
        new_data = np.array(new_data)
    
    # 数据标准化（这里需要加载训练时使用的scaler）
    # 在实际应用中，你应该保存和加载scaler
    scaler = StandardScaler()
    scaler.fit(new_data)  # 这里只是示例，实际应该使用训练时的scaler
    
    new_data_scaled = scaler.transform(new_data)
    
    # 创建序列
    if len(new_data_scaled) >= sequence_length:
        sequences, _ = create_sequences(new_data_scaled, sequence_length)
        sequences = torch.FloatTensor(sequences).to(device)
        
        # 进行预测
        with torch.no_grad():
            outputs = model(sequences)
            _, predicted = torch.max(outputs.data, 1)
        
        return predicted.cpu().numpy(), sequences.cpu().numpy()
    else:
        print("数据长度不足以创建序列")
        return None, None

def predict_single_sequence(model_path, sequence, device):
    """预测单个序列"""
    
    # 加载模型
    model = load_model(create_model, model_path, device)
    model.eval()
    
    # 确保序列是合适的形状
    if isinstance(sequence, list):
        sequence = np.array(sequence)
    
    if len(sequence.shape) == 1:
        sequence = sequence.reshape(1, -1, 1)  # (batch_size, sequence_length, input_size)
    elif len(sequence.shape) == 2:
        sequence = sequence.reshape(1, sequence.shape[0], sequence.shape[1])
    
    sequence = torch.FloatTensor(sequence).to(device)
    
    # 进行预测
    with torch.no_grad():
        outputs = model(sequence)
        probabilities = torch.softmax(outputs, dim=1)
        _, predicted = torch.max(outputs.data, 1)
    
    return predicted.cpu().numpy(), probabilities.cpu().numpy()

def demo_prediction():
    """演示预测功能"""
    
    device = torch.device(TRAIN_CONFIG['device'])
    model_path = SAVE_CONFIG['model_save_path']
    
    # 检查模型是否存在
    import os
    if not os.path.exists(model_path):
        print(f"模型文件 {model_path} 不存在，请先运行 train.py 进行训练")
        return
    
    print("正在加载模型进行预测演示...")
    
    # 创建一些测试数据
    test_sequence = np.random.randn(TRAIN_CONFIG['sequence_length'], MODEL_CONFIG['input_size'])
    
    # 进行预测
    predicted, probabilities = predict_single_sequence(model_path, test_sequence, device)
    
    if predicted is not None:
        print(f"测试序列预测结果:")
        print(f"预测类别: {predicted[0]}")
        print(f"各类别概率: {probabilities[0]}")
        
        # 如果有真实数据，可以进行更详细的演示
        try:
            # 加载真实数据进行演示
            df = pd.read_csv(DATA_CONFIG['data_path'])
            features = df[DATA_CONFIG['feature_columns']].values
            
            # 使用最后几个数据点进行预测
            if len(features) >= TRAIN_CONFIG['sequence_length']:
                last_sequence = features[-TRAIN_CONFIG['sequence_length']:]
                predicted_real, _ = predict_single_sequence(model_path, last_sequence, device)
                
                print(f"\n使用真实数据最后{len(last_sequence)}个点进行预测:")
                print(f"数据序列: {last_sequence.flatten()}")
                print(f"预测类别: {predicted_real[0]}")
                
                # 显示真实标签（如果可用）
                if len(df) > TRAIN_CONFIG['sequence_length']:
                    true_label = df[DATA_CONFIG['target_column']].iloc[-1]
                    print(f"真实标签: {true_label}")
                    print(f"预测是否正确: {predicted_real[0] == true_label}")
        
        except Exception as e:
            print(f"加载真实数据时出错: {e}")
    
    print("\n预测演示完成!")

if __name__ == "__main__":
    demo_prediction()